"""
YOLOv8 完整模型调优脚本
功能：
1. 超参数网格搜索
2. 数据增强配置
3. 学习率策略调整
4. 模型结构微调
5. 自动保存最佳配置
"""
import os
import yaml
import torch
from pathlib import Path
from ultralytics import YOLO
import logging
import numpy as np
from sklearn.model_selection import ParameterGrid

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[logging.FileHandler('logs/advanced.log'), logging.StreamHandler()]
)
logger = logging.getLogger(__name__)


class YOLOv8Tuner:
    def __init__(self, config_path="configs/default.yaml"):
        self.config = self.load_config(config_path)
        self.setup_paths()
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.best_metrics = {
            'map50': 0,
            'params': {},
            'model_path': ''
        }

        self.resume_from = self.config['tuning']['resume_from']

    def load_config(self, config_path):
        """加载调优配置"""
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)

        # 生成参数网格
        self.param_grid = {
            'lr0': np.linspace(config['tuning']['lr_range'][0],
                               config['tuning']['lr_range'][1],
                               config['tuning']['lr_range'][2]).tolist(),
            **config['tuning']['search_space']
        }
        return config

    def setup_paths(self):
        """设置路径"""
        base_dir = Path(__file__).parent
        self.dataset_yaml = str(base_dir / self.config['paths']['processed']['dataset'])
        self.project_dir = str(base_dir / self.config['paths']['runs']['tune'])
        self.best_config_dir = base_dir / "configs" / "best_configs"
        self.best_config_dir.mkdir(parents=True, exist_ok=True)

    def train_model(self, params):
        """使用给定参数训练模型"""
        try:

            if self.resume_from:
                model = YOLO(self.config['tuning']["path"])
                logger.info(f"从现有模型继续训练: {self.config['tuning']['path']}")
            else:
                model = YOLO(self.config['training']['base']['model'])

            # 合并固定参数和搜索参数
            train_args = {
                'data': self.dataset_yaml,
                'epochs': self.config['training']['base']['epochs'],
                'imgsz': self.config['training']['base']['imgsz'],
                'batch': self.config['training']['base']['batch'],
                'device': self.device,
                'project': self.project_dir,
                'name': f"trial_{len(os.listdir(self.project_dir)) if os.path.exists(self.project_dir) else 0}",
                'exist_ok': True,
                **params,
                **self.config['fixed_params']
            }

            # 训练
            results = model.train(**train_args)
            return results
        except Exception as e:
            logger.error(f"训练失败: {str(e)}")
            return None

    def evaluate_model(self, model_path):
        """评估模型性能"""
        model = YOLO(model_path)
        metrics = model.val(
            data=self.dataset_yaml,
            split='val',
            device=self.device,
            plots=False,
            save_json=True
        )
        return {
            'map50': metrics.box.map50,
            'map': metrics.box.map,
            'precision': metrics.box.p.mean(),
            'recall': metrics.box.r.mean()
        }

    def run_tuning(self):
        """执行参数搜索"""
        param_combinations = list(ParameterGrid(self.param_grid))
        logger.info(f"开始超参数搜索，共 {len(param_combinations)} 种组合")

        for i, params in enumerate(param_combinations):
            logger.info(f"\n=== 尝试组合 {i + 1}/{len(param_combinations)} ===")
            logger.info(f"参数: {params}")
            logger.info(yaml.dump(params, allow_unicode=True))

            # 训练并评估
            results = self.train_model(params)
            if not results:
                continue

            metrics  = self.evaluate_model(Path(results.save_dir) / 'weights' / 'best.pt')
            logger.info(f"评估指标: mAP50={metrics['map50']:.4f} (当前最佳: {self.best_metrics['map50']:.4f})")

            # 更新最佳结果
            if metrics['map50'] > self.best_metrics['map50']:
                self.best_metrics = {
                    'map50': metrics['map50'],
                    'params': params,
                    'model_path': str(Path(results.save_dir) / 'weights' / 'best.pt'),
                    'full_metrics': metrics
                }
                self.save_best_config()

    def save_best_config(self):
        """保存最佳配置"""
        config_path = self.best_config_dir / f"best_map50_{self.best_metrics['map50']:.4f}.yaml"

        config = {
            'performance': self.best_metrics['full_metrics'],
            'parameters': self.best_metrics['params'],
            'model_path': self.best_metrics['model_path'],
            'dataset': self.dataset_yaml
        }

        with open(config_path, 'w') as f:
            yaml.dump(config, f)

        logger.info(f"保存最佳配置到: {config_path}")

    def final_report(self):
        """生成最终报告"""
        logger.info("\n" + "=" * 50)
        logger.info("=== 调优结果 ===")
        logger.info(f"最佳 mAP50: {self.best_metrics['map50']:.4f}")
        logger.info("最佳参数组合:")
        logger.info(yaml.dump(self.best_metrics['params'], allow_unicode=True))
        logger.info(f"模型路径: {self.best_metrics['model_path']}")
        logger.info("=" * 50)


if __name__ == "__main__":
    try:
        logger.info("=" * 50)
        logger.info("YOLOv8 超参数调优")
        logger.info("=" * 50)

        tuner = YOLOv8Tuner()
        tuner.run_tuning()
        tuner.final_report()
    except Exception as e:
        logger.error(f"调优失败: {str(e)}")
        exit(1)
